// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/common/memory/memory.h>
#include <pollux/vector/base_vector.h>

/// These 2 definitions should be included by user from either
///   1. <arrow/c/abi.h> or
///   2. "pollux/vector/arrow/Abi.h"
struct ArrowArray;
struct ArrowSchema;

enum class TimestampUnit : uint8_t {
    kSecond = 0 /*10^0 second is equal to 1 second*/,
    kMilli = 3 /*10^3 milliseconds are equal to 1 second*/,
    kMicro = 6 /*10^6 microseconds are equal to 1 second*/,
    kNano = 9 /*10^9 nanoseconds are equal to 1 second*/
};

struct ArrowOptions {
    bool flattenDictionary{false};
    // NOTE: flattenConstant is only supported for scalar types.
    bool flattenConstant{false};
    TimestampUnit timestampUnit = TimestampUnit::kNano;
    std::optional<std::string> timestampTimeZone{std::nullopt};
    // Export VARCHAR and VARBINARY to Arrow 15 StringView format
    bool exportToStringView = false;
};

namespace kumo::pollux {
    /// Export a generic Pollux Vector to an ArrowArray, as defined by Arrow's C data
    /// interface:
    ///
    ///   https://arrow.apache.org/docs/format/CDataInterface.html
    ///
    /// The output ArrowArray needs to be allocated by the consumer (either in the
    /// heap or stack), and after usage, the standard REQUIRES the client to call
    /// the release() function (or memory will leak).
    ///
    /// After exporting, the ArrowArray will hold ownership to the underlying Vector
    /// being referenced, so the consumer does not need to explicitly hold on to the
    /// input Vector shared_ptr.
    ///
    /// The function takes a memory pool where allocations will be made (in cases
    /// where the conversion is not zero-copy, e.g. for strings) and throws in case
    /// the conversion is not implemented yet.
    ///
    /// Example usage:
    ///
    ///   ArrowArray arrowArray;
    ///   arrow::export_to_arrow(inputVector, arrowArray);
    ///   inputVector.reset(); // don't need to hold on to this shared_ptr.
    ///
    ///   (use arrowArray)
    ///
    ///   arrowArray.release(&arrowArray);
    void export_to_arrow(
        const VectorPtr &vector,
        ArrowArray &arrowArray,
        memory::MemoryPool *pool,
        const ArrowOptions &options = ArrowOptions{});

    /// Export the type of a Pollux vector to an ArrowSchema.
    ///
    /// The guidelines on API usage and memory management are the same as the ones
    /// described for the VectorPtr->ArrowArray export function.
    ///
    /// The function throws in case there was no valid conversion available.
    ///
    /// Example usage:
    ///
    ///   ArrowSchema arrowSchema;
    ///   arrow::export_to_arrow(inputType, arrowSchema);
    ///   inputType.reset(); // don't need to hold on to this shared_ptr.
    ///
    ///   (use arrowSchema)
    ///
    ///   arrowSchema.release(&arrowSchema);
    ///
    /// NOTE: Since Arrow couples type and encoding, we need both Pollux type and
    /// actual data (containing encoding) to create an ArrowSchema.
    void export_to_arrow(
        const VectorPtr &,
        ArrowSchema &,
        const ArrowOptions & = ArrowOptions{});

    /// Import an ArrowSchema into a Pollux Type object.
    ///
    /// This function does the exact opposite of the function above. TypePtr carries
    /// all buffers they need to represent types, so after this function returns,
    /// the client is free to release any buffers associated with the input
    /// ArrowSchema object.
    ///
    /// The function throws in case there was no valid conversion available.
    ///
    /// Example usage:
    ///
    ///   ArrowSchema arrowSchema;
    ///   ... // fills arrowSchema
    ///   auto type = arrow::import_from_arrow(arrowSchema);
    ///
    ///   arrowSchema.release(&arrowSchema);
    ///
    TypePtr import_from_arrow(const ArrowSchema &arrowSchema);

    /// Import an ArrowArray and ArrowSchema into a Pollux vector.
    ///
    /// This function takes both an ArrowArray (which contains the buffers) and an
    /// ArrowSchema (which describes the data type), since a Pollux vector needs
    /// both (buffer and type). A memory pool is also required, since all vectors
    /// carry a pointer to it, but not really used in most cases - unless the
    /// conversion itself requires a new allocation. In most cases no new
    /// allocations are required, unless for arrays of varchars (or varbinaries) and
    /// complex types written out of order.
    ///
    /// The new Pollux vector returned contains only references to the underlying
    /// buffers, so it's the client's responsibility to ensure the buffer's
    /// lifetime.
    ///
    /// The function throws in case the conversion fails.
    ///
    /// Example usage:
    ///
    ///   ArrowSchema arrowSchema;
    ///   ArrowArray arrowArray;
    ///   ... // fills structures
    ///   auto vector = arrow::import_from_arrow(arrowSchema, arrowArray, pool);
    ///   ... // ensure buffers in arrowArray remain alive while vector is used.
    ///
    VectorPtr import_from_arrow_as_viewer(
        const ArrowSchema &arrowSchema,
        const ArrowArray &arrowArray,
        memory::MemoryPool *pool);

    /// Import an ArrowArray and ArrowSchema into a Pollux vector, acquiring
    /// ownership over the input data.
    ///
    /// Similar to import_from_arrow_as_viewer but the ownership of arrowSchema and
    /// arrowArray will be taken over. Specifically, the returned Vector will own a
    /// copy of arrowSchema and arrowArray.
    ///
    /// The inputs arrowSchema and arrowArray will be marked as released by setting
    /// their release callback to nullptr
    /// (https://arrow.apache.org/docs/format/CDataInterface.html). Afterwards, the
    /// returned Vector will be responsible for calling the release callbacks when
    /// destructed.
    VectorPtr import_from_arrow_as_owner(
        ArrowSchema &arrowSchema,
        ArrowArray &arrowArray,
        memory::MemoryPool *pool);
} // namespace kumo::pollux
